from insightface.app import FaceAnalysis
import cv2
import numpy as np
from pymongo import MongoClient
from datetime import datetime
import os
from typing import List, Dict, Union
import logging

class FaceSearcher:
    def __init__(self, mongo_uri: str = "mongodb://192.168.0.234:27017/"):
        """
        初始化人脸搜索器
        
        Args:
            mongo_uri: MongoDB连接字符串
        """
        # 配置日志
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)
        
        # 初始化InsightFace
        self.app = FaceAnalysis(
            providers=['CPUExecutionProvider'],  # 移除CUDA，因为您的系统不支持
            allowed_modules=['detection', 'recognition']
        )
        self.app.prepare(ctx_id=0, det_size=(640, 640))
        
        # 连接MongoDB
        try:
            # 添加连接超时设置
            self.client = MongoClient(
                mongo_uri,
                serverSelectionTimeoutMS=5000,  # 5秒连接超时
                connectTimeoutMS=5000
            )
            
            # 测试连接
            self.client.server_info()
            self.logger.info("MongoDB连接测试成功")
            
            self.db = self.client.az_face
            self.face_collection = self.db.face_vectors
            
            # 创建索引
            try:
                self.face_collection.create_index([("face_id", 1)], unique=True)
                self.face_collection.create_index([("person_id", 1)])
                self.logger.info("MongoDB索引创建成功")
            except Exception as e:
                self.logger.warning(f"创建索引时出错: {str(e)}")
            
        except Exception as e:
            self.logger.error(f"MongoDB连接失败: {str(e)}")
            raise ConnectionError(f"无法连接到MongoDB服务器 {mongo_uri}: {str(e)}")
    
    def _extract_face_embedding(self, image_path: str) -> Union[np.ndarray, None]:
        """
        从图片中提取人脸特征向量
        
        Args:
            image_path: 图片路径
            
        Returns:
            人脸特征向量或None（如果没有检测到人脸）
        """
        try:
            img = cv2.imread(image_path)
            if img is None:
                self.logger.error(f"无法读取图片: {image_path}")
                return None
                
            faces = self.app.get(img)
            
            if not faces:
                self.logger.warning(f"未检测到人脸: {image_path}")
                return None
                
            # 返回检测到的第一个人脸的特征向量
            return faces[0].embedding
            
        except Exception as e:
            self.logger.error(f"特征提取失败: {str(e)}")
            return None
    
    def add_face(self, image_path: str, person_id: str, metadata: Dict = None) -> bool:
        """
        添加人脸到数据库
        
        Args:
            image_path: 图片路径
            person_id: 人物ID
            metadata: 额外的元数据信息
            
        Returns:
            bool: 是否添加成功
        """
        embedding = self._extract_face_embedding(image_path)
        if embedding is None:
            return False
            
        try:
            face_document = {
                "face_id": f"face_{datetime.now().timestamp()}",
                "person_id": person_id,
                "embedding": embedding.tolist(),
                "created_at": datetime.now(),
                "image_path": image_path,
                "metadata": metadata or {}
            }
            
            self.face_collection.insert_one(face_document)
            self.logger.info(f"成功添加人脸: {person_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"添加人脸失败: {str(e)}")
            return False
    
    def batch_add_faces(self, image_dir: str, person_id: str) -> int:
        """
        批量添加某个人的人脸照片
        
        Args:
            image_dir: 包含人脸图片的目录
            person_id: 人物ID
            
        Returns:
            int: 成功添加的人脸数量
        """
        success_count = 0
        for filename in os.listdir(image_dir):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
                image_path = os.path.join(image_dir, filename)
                if self.add_face(image_path, person_id):
                    success_count += 1
                    
        return success_count
    
    def search_similar_faces(self, 
                           query_image: str, 
                           threshold: float = 0.5,
                           top_k: int = 5) -> List[Dict]:
        """
        搜索相似人脸
        
        Args:
            query_image: 查询图片路径
            threshold: 相似度阈值
            top_k: 返回的最相似人脸数量
            
        Returns:
            List[Dict]: 相似人脸列表，包含相似度分数和人物信息
        """
        query_embedding = self._extract_face_embedding(query_image)
        if query_embedding is None:
            return []
            
        results = []
        # 从数据库获取所有人脸特征向量
        for face in self.face_collection.find():
            db_embedding = np.array(face['embedding'])
            # 计算余弦相似度
            similarity = np.dot(query_embedding, db_embedding) / \
                        (np.linalg.norm(query_embedding) * np.linalg.norm(db_embedding))
            
            if similarity >= threshold:
                results.append({
                    'person_id': face['person_id'],
                    'face_id': face['face_id'],
                    'similarity': float(similarity),
                    'image_path': face['image_path'],
                    'metadata': face['metadata']
                })
        
        # 按相似度排序
        results.sort(key=lambda x: x['similarity'], reverse=True)
        return results[:top_k]
    
    def batch_search(self, 
                    query_images: List[str], 
                    threshold: float = 0.5,
                    top_k: int = 5) -> Dict[str, List[Dict]]:
        """
        批量搜索多张图片
        
        Args:
            query_images: 查询图片路径列表
            threshold: 相似度阈值
            top_k: 每张图片返回的最相似人脸数量
            
        Returns:
            Dict[str, List[Dict]]: 每张查询图片的搜索结果
        """
        results = {}
        for image_path in query_images:
            results[image_path] = self.search_similar_faces(
                image_path, 
                threshold=threshold,
                top_k=top_k
            )
        return results
    
    def delete_person(self, person_id: str) -> int:
        """
        删除某个人的所有人脸记录
        
        Args:
            person_id: 人物ID
            
        Returns:
            int: 删除的记录数量
        """
        result = self.face_collection.delete_many({"person_id": person_id})
        return result.deleted_count
    
    def __del__(self):
        """清理资源"""
        if hasattr(self, 'client'):
            self.client.close() 